using SharpToken;

namespace Aurora.AI.Common.Helper;

/// <summary>
/// token helper
/// </summary>
public static class TokenHelper
{
    private static GptEncoding _encoding;

    /// <summary>
    ///     获取GptEncoding
    /// </summary>
    /// <returns></returns>
    public static GptEncoding GetGptEncoding()
    {
        _encoding ??= GptEncoding.GetEncodingForModel("gpt-4");

        return _encoding;
    }

    /// <summary>
    ///     计算token
    /// </summary>
    /// <param name="tokens"></param>
    /// <returns></returns>
    public static int ComputeToken(params string[] tokens)
    {
        return tokens.Where(token => !string.IsNullOrEmpty(token)).Sum(token => GetGptEncoding().Encode(token).Count);
    }

    /// <summary>
    /// 
    /// </summary>
    /// <param name="content"></param>
    /// <returns></returns>
    public static int GetTotalTokens(params string[] content)
    {
        return content.Sum(item => GetGptEncoding().Encode(item).Count);
    }

    public static int GetTokens(string content)
    {
        return GetGptEncoding().Encode(content).Count;
    }
}